# %%
from conn import ac_task_col, clip_annotated_record, check_frame_label_col, frame_col
from collections import defaultdict
from tqdm import tqdm

# %%

query = dict(accept_status=20)
projection = dict(clip_id=1, accept_status=1, annotate_type=1)
total = clip_annotated_record.count_documents(query)
print(total)


# %%
groups = defaultdict(list)
for clip_id in tqdm(clip_annotated_record.find(query, projection), total=total):
    groups[clip_id["annotate_type"]].append(clip_id["clip_id"])


# %%

for k, v in groups.items():
    print(k, len(v))

# %%
clip_ids = groups[31]
for clip_id in tqdm(clip_ids):
    if ac_task_col.find_one({"clip_ids": clip_id, "annotate_type": 31}) is None:
        continue
    frame_cnt = frame_col.count_documents({"clip_id": clip_id}) // 10  # 1hz
    label_frame_cnt = check_frame_label_col.count_documents(
        {"clip_id": clip_id, "annotate_type": 31}
    )
    if label_frame_cnt < frame_cnt:
        print(clip_id, label_frame_cnt, frame_cnt)

# %%
groups = frame_col.aggregate(
    [{"$group": {"_id": "$clip_id", "cnt": {"$sum": 1}}}], allowDiskUse=True
).to_list()

frame_cnt_set = set([g["cnt"] for g in groups])

# %%
print(frame_cnt_set)
